# 导入必要的库
from datasets import load_dataset  # 用于加载数据集
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer  # Hugging Face 的核心组件
from peft import LoraConfig, TaskType, get_peft_model,PeftModel  # 用于高效微调的 PEFT 库
import torch  # PyTorch 框架
import os  # 操作系统相关功能

# 设置环境变量以减少显存碎片
# 通过设置 expandable_segments:True 优化 PyTorch 的 CUDA 内存分配
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"

# 清理显存
# 释放 GPU 上可能残留的未使用显存
torch.cuda.empty_cache()

# 加载数据集
# 从指定路径加载 Arrow 格式的中文 Alpaca 数据集，仅使用训练集
ds = load_dataset(
    "arrow",
    data_files="./fengjay01___alpaca_data_zh_51k/default/0.0.0/3ddd31c70ca8b6ec10fcfe8d907aaa5dfc6111f7/alpaca_data_zh_51k-train.arrow"
)
ds = ds['train']

# 加载 tokenizer
# 从预训练模型路径加载 LLaMA-3.1-8B-Instruct 的分词器，并设置 pad_token 为 eos_token
tokenizer = AutoTokenizer.from_pretrained("./model/ZhipuAI/glm-4-9b/",trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token

query = "如何提高学习效率"
inputs = tokenizer.apply_chat_template(
    [{"role": "user", "content": query}],  # 构造用户角色的消息
    add_generation_prompt=True,  # 添加生成提示符
    tokenize=True,  # 进行分词
    return_tensors="pt",  # 返回PyTorch张量
    return_dict=True  # 返回字典格式
)

# 将输入移动到CUDA设备
inputs = inputs.to("cuda")
# print(inputs)
print("需要拼接的格式为:",tokenizer.decode(inputs["input_ids"][0]))
# 数据预处理函数
# 将数据集中的样本转换为模型输入格式，包括 input_ids、attention_mask 和 labels
def process_func(example):
    """
    处理单个样本，生成输入ID、注意力掩码和标签，用于模型训练或推理

    Args:
        example (dict): 包含"instruction", "input", 和"output"的样本字典

    Returns:
        dict: 包含"input_ids", "attention_mask", 和"labels"的字典
    """
    MAX_LENGTH = 256
    input_ids, attention_mask, labels = [], [], []

    # 获取纯文本的instruction，不拼接"Human:"前缀
    instruction = "\n".join([example["instruction"], example["input"]]).strip()

    # 应用聊天模板，将instruction转换为模型输入
    instruction = tokenizer.apply_chat_template(
        [{"role": "user", "content": instruction}],  # 构造用户角色的消息
        add_generation_prompt=True,  # 添加生成提示符
        tokenize=True,  # 进行分词
        return_tensors="pt",  # 返回PyTorch张量
        return_dict=True  # 返回字典格式
    )

    # 处理response，添加换行符并禁用特殊标记
    response = tokenizer("\n" + example["output"], add_special_tokens=False)

    # 合并输入ID、注意力掩码和标签，转换为列表并添加结束标记
    input_ids = instruction["input_ids"][0].numpy().tolist() + response["input_ids"] + [tokenizer.eos_token_id]
    attention_mask = instruction["attention_mask"][0].numpy().tolist() + response["attention_mask"] + [1]
    labels = [-100] * len(instruction["input_ids"][0].numpy().tolist()) + response["input_ids"] + [tokenizer.eos_token_id]

    # 截断到最大长度
    if len(input_ids) > MAX_LENGTH:
        input_ids = input_ids[:MAX_LENGTH]
        attention_mask = attention_mask[:MAX_LENGTH]
        labels = labels[:MAX_LENGTH]

    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels
    }

# 应用预处理函数到数据集并移除原始列
tokenized_ds = ds.map(process_func, remove_columns=ds.column_names)

# 加载预训练模型

model = AutoModelForCausalLM.from_pretrained(
    "./model/ZhipuAI/glm-4-9b/",
    low_cpu_mem_usage=True,  # 减少 CPU 内存占用
    torch_dtype=torch.half,# 使用半精度浮点数
    trust_remote_code=True
)
model.gradient_checkpointing_enable()  # 启用梯度检查点以减少显存占用
model = model.cuda()  # 将模型移动到 GPU

# 配置 LoRA（低秩适配）
# 使用 LoraConfig 设置 LoRA 参数，当前仅指定任务类型，注释部分为更详细的配置示例
# config = LoraConfig(
#     task_type=TaskType.CAUSAL_LM,  # 任务类型为因果语言模型
#     r=8,  # LoRA 秩
#     lora_alpha=16,  # LoRA 缩放因子
#     lora_dropout=0.1,  # LoRA 丢弃率
#     target_modules=["q_proj", "v_proj"]  # 应用 LoRA 的模块
# )
config = LoraConfig(task_type=TaskType.CAUSAL_LM,target_modules=["query_key_value"])
config
model = get_peft_model(model, config)  # 将 LoRA 应用到模型
model.print_trainable_parameters()  # 打印可训练参数数量

# 配置训练参数
# 使用 TrainingArguments 定义训练超参数
args = TrainingArguments(
    output_dir="./output",  # 模型检查点和输出的保存目录
    per_device_train_batch_size=8,  # 每个设备的训练批次大小
    num_train_epochs=3,  # 训练的总轮数
    logging_dir="./logs",  # 日志保存目录
    logging_steps=100,  # 每 100 步记录一次日志
    save_steps=500,  # 每 500 步保存一次模型检查点
    save_total_limit=2,  # 最多保留 2 个检查点
)

# 初始化 Trainer
# 配置 Trainer 用于训练模型，指定模型、参数、数据集和数据整理器
trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenized_ds.select(range(1000)),  # 使用前 1000 个样本进行测试
    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True)  # 数据整理器，自动填充序列
)

 # 开始训练
 # 调用 trainer.train() 启动模型训练（当前被注释）
trainer.train()

# 模型推理
# 构造推理输入，格式为 "Human: 指令"，并移动到模型所在的设备
ipt = tokenizer(
    "Human:{}\n{}".format("如何提高学习效率", "").strip() + "\n\nAssistant: ",
    return_tensors='pt'
).to(model.device)
# 使用模型生成响应，设置最大长度为 384，并启用采样
print(tokenizer.decode(
    model.generate(**ipt, max_length=384, do_sample=True)[0],
    skip_special_tokens=True
))